// file: drive/keyboard.c
// autor: jiang xinpeng
// time:2021.4.4
// copyright: (C) by jiangxinpeng,All right are reserved.

#include <os/driver.h>
#include <os/initcall.h>
#include <os/fifoio.h>
#include <os/hardirq.h>
#include <os/debug.h>
#include <arch/interrupt.h>
#include <driver/keyboard.h>
#include <sys/res.h>
#include <sys/ioctl.h>
#include <lib/list.h>

//#define DEBUG_KEYBOARD
 
static uint32_t ScanCodeToEventCode(uint32_t key);
static int InitController();
static int TestController();
static void OpenFirstPort();
static void OpenSecondPort();
static void SendCmd(uint8_t cmd);
static uint8_t CheckDeviceType(uint8_t port);
static int CheckOutBuffer();
static uint8_t ReadData();
static void WriteData(uint8_t data);
static void WriteDataTo(uint8_t data, uint8_t port);
static int CheckInBuffer();
static void WaitInBuffer();
static void WaitOutBuffer();
static inline int TestPort(uint8_t id);
static inline void CheckAck();

static int KeyboardHandler(irqno_t irq, void *data);
static int KeyboardProbe(device_extension_t *extension, uint8_t id);
static void KeyboardThread(void *arg);
static void KeyboardSetLed(device_extension_t *extension);

static uint8_t CheckDeviceType(uint8_t port)
{
    uint8_t data;
    uint8_t identify0;
    uint8_t identify1;
    uint8_t type = PS2_TYPE_NONE;

    // send cmd get identify from device
    WriteDataTo(PS2_DEVICE_CMD_IDENTIFY, port);
    // wait device return ACK
    data = ReadData();
    if (data == PS2_DEVICE_CMD_RETURN_ACK)
    {
        // receive identify byte from device
        int timeout = 1000;
        // identify 0
        while (timeout-- && CheckInBuffer())
            ;
        if (timeout == 0)
        {
            identify0 = 0xff;
        }
        else
        {
            identify0 = In8(PS2CTRL_IO_DATA_PORT);
        }
        // identify1
        timeout = 1000;
        while (timeout-- && CheckInBuffer())
            ;
        if (timeout == 0)
        {
            identify1 = 0xff;
        }
        else
        {
            identify1 = In8(PS2CTRL_IO_DATA_PORT);
        }

        // if identify is any mouse type
        if (identify0 == 0x00 || identify0 == 0x03 || identify0 == 0x04)
        {
            type = PS2_TYPE_MOUSE;
        }
        else
        {
            // check if is keyboard type device
            if ((identify0 == 0xAB && identify1 == 0x41) || identify0 == 0xAB && identify1 == 0xC1 || identify0 == 0xAB && identify1 == 0x83)
                type = PS2_TYPE_KEYBOARD;
            else
                type = PS2_TYPE_KEYBOARD;
        }
    }
    return type;
}

// read data from ps2 controller data port
static uint8_t ReadData()
{
    uint8_t data;

    // check output buffer
    WaitOutBuffer();
    // read data
    data = In8(PS2CTRL_IO_DATA_PORT);
    return data;
}

static void SendCmd(uint8_t cmd)
{
    // wait input buffer empty
    WaitInBuffer();
    // send cmd
    Out8(PS2CTRL_IO_CMD_PORT, cmd);
}

static void WaitInBuffer()
{
    uint8_t data;
    while ((data = CheckInBuffer()))
        ;
}

static void WaitOutBuffer()
{
    uint8_t data;
    while (!(data = CheckOutBuffer()))
        ;
}

static void WriteData(uint8_t data)
{
    // wait input buffer empty
    WaitInBuffer();
    // send data to second ps/2 port
    Out8(PS2CTRL_IO_DATA_PORT, data);
}

static void WriteDataTo(uint8_t data, uint8_t port)
{
    uint8_t flag, count;
    // first ps/2 port
    if (port == 0x00)
    {
        WaitInBuffer();
        // send data to data port
        Out8(PS2CTRL_IO_DATA_PORT, data);
    }
    // second ps/2 port
    // we need first send a cmd to enable second port
    else
    {
        if (port = 0x01)
        {
            // send cmd to write to second port
            SendCmd(PS2CTRL_CMD_WRITENEXTBYTE_TO_SECONDPORTINBUFF);
            WaitInBuffer();
            // send data to second ps/2 port
            Out8(PS2CTRL_IO_DATA_PORT, data);
        }
    }
}

// return: 0 empty
//         1 full
static int CheckInBuffer()
{
    uint8_t data;
    // read status
    data = In8(PS2CTRL_IO_STATUS_PORT);
    if (data & PS2CTRL_STATUS_INFULL)
    {
        return 1;
    }
    return 0;
}

void FlushBuffer()
{
    uint8_t data;
    // if output buffer no empty
    while (CheckOutBuffer())
    {
        data = In8(PS2CTRL_IO_DATA_PORT);
    }
}

uint8_t ReadConfig()
{
    uint8_t config;

    // send read config bytes cmd
    SendCmd(PS2CTRL_CMD_READCONFIGBYTE);
    // wait and receive data
    config = ReadData();
    return config;
}

void WriteConfig(uint8_t config)
{
    // send write config space cmd
    SendCmd(PS2CTRL_CMD_WRITECONFIGBYTE);
    // write config data to data port
    WriteData(config);
}

static uint32_t ScanCodeToEventCode(uint32_t key)
{
    int i;
    for (i = 0; i < ARRAY_SIZE(map_table); i++)
    {
        if (map_table[i] == key)
        {
            return map_table[i + 1];
        }
    }
    return KEY_UNKNOWN;
}

static iostatus_t KeyboardEnter(driver_object_t *driver)
{
    iostatus_t status = IO_SUCCESS;

    device_object_t *devobj;
    device_extension_t *extension;
    void *buff;
    uint8_t err = 0;

    FlushBuffer(); //flush buffer

    // probe ps2 keyboard and create device
    for (int i = 0; i < 2; i++)
    {
        // create device
        status = IoCreateDevice(driver, sizeof(device_extension_t), DEVICE_NAME, DEVICE_TYPE_KEYBOARD, &devobj);
        if (status != IO_SUCCESS)
        {
            KPrint(PRINT_ERR "KeyboardEnter: device %s object create failed!\n", DEVICE_NAME);
            return status;
        }
        devobj->flags = 0;
        extension = devobj->device_extension;

        if (KeyboardProbe(extension, i) < 0) // probe device info
        {
            err++;
            status = IO_FAILED;
            IoDeleteDevice(devobj);
            continue;
        }

        // set extension info
        extension->devobj = devobj;
        extension->l_shift = extension->r_shift = extension->l_ctrl = extension->r_ctrl = 0;
        extension->capslock = 0;
        extension->numlock = 0;
        extension->scrolllock = 0;
        extension->open = 0;
        extension->exist = 1;
        extension->irq = IRQ1_KEYBOARD;

        // init fifo and event buff
        buff = KMemAlloc(DEVICE_BUFFER_SIZE);
        if (!buff)
        {
            KPrint("[keyboard] buffer alloc failed!\n");
            IoDeleteDevice(devobj);
            return IO_FAILED;
        }
        FifoIoInit(&extension->fifoio, buff, DEVICE_BUFFER_SIZE);
        InputEventInit(&extension->eventbuff);

        // start kernel thread
        TaskCreate("keyboard", TASK_PRIOR_LEVEL_HIGH, KeyboardThread, extension);

        // register irq
        IrqRegister(extension->irq, KeyboardHandler, IRQ_DISABLE, "IRQ1", DEVICE_NAME, extension);
    }

    if (err > 1)
    {
        KPrint("%s: keyboard device probe failed! this system no found ps2 keyboard!\n", __func__);
        status = IO_FAILED;
    }
    else
    {
        status = IO_SUCCESS;
    }
    return status;
}

static iostatus_t KeyboardExit(driver_object_t *driver)
{
    device_object_t *device, *next;
    list_traversal_all_owner_to_next_safe(device, next, &driver->device_list, list)
    {
        IoDeleteDevice(device);
    }
    string_del(&driver->name);
    return IO_SUCCESS;
}

static iostatus_t KeyboardOpen(device_object_t *devobj, io_request_t *ioreq)
{
    device_extension_t *extension = devobj->device_extension;
    iostatus_t status = IO_SUCCESS;

    extension->open = 1;
    ioreq->io_status.status = IO_SUCCESS;
    ioreq->io_status.info = 0;
    IoCompleteRequest(ioreq);
    return status;
}

// read bytes from keyboard buffer
static inline uint8_t GetBytesFromBuff(device_extension_t *extension)
{
    return FifoIoGet(&extension->fifoio);
}

// read keys and deal with
static uint32_t KeyboardDoRead(device_extension_t *extension)
{
    uint8_t scancode = 0;
    uint32_t key = 0, make = 0;
    uint32_t *keyrow = NULL;
    uint8_t is_pausebreak = 0;

    extension->code_with_e0 = 0;

    scancode = GetBytesFromBuff(extension);

    if (scancode == 0)
        return 0;

    if (scancode == 0xe1) // pause/break keydown
    {
        uint8_t pausebreak_code[] = {0xE1, 0x1D, 0x45, 0xE1, 0x9D, 0xC5};
        for (int i = 1; i < 6; i++)
        {
            if (GetBytesFromBuff(extension) != pausebreak_code[i])
            {
                is_pausebreak = 1;
                break;
            }
        }
        if (is_pausebreak)
            key = KEYBOARD_PAUSEBREAK;
    }
    else
    {
        // printscreen keydown
        if (scancode == 0xe0)
        {
            scancode = GetBytesFromBuff(extension);

            if (scancode == 0x2A)
            {
                if (GetBytesFromBuff(extension) == 0xE0)
                {
                    if (GetBytesFromBuff(extension) == 0x37)
                    {
                        key = KEYBOARD_PRINTSCREEN;
                        make = 1;
                        KPrint("[keyboard] printscreen key down\n");
                    }
                }
            }

            // printscreen keyup
            if (scancode == 0xB7)
            {
                if (GetBytesFromBuff(extension) == 0xE0)
                {
                    if (GetBytesFromBuff(extension) == 0xAA)
                    {
                        key = KEYBOARD_PRINTSCREEN;
                        make = 0;
                        KPrint("[keyboard] printscreen key down\n");
                    }
                }
            }

            if (key == 0) // no printscreen
            {
                extension->code_with_e0 = 1;
            }
        }
    }

    if ((key != KEYBOARD_PAUSEBREAK) && (key != KEYBOARD_PRINTSCREEN))
    {
        make = ((scancode & KEYBOARD_FLAG_BREAK_MASK) ? 0 : 1);

        keyrow = &(kbd_keymap[(scancode & 0x7F) * KEYMAP_COLS]);

        extension->colume = 0;
        // check caps
        int caps = extension->l_shift || extension->r_shift;
        if (extension->capslock)
        {
            if (keyrow[0] >= 'a' && keyrow[0] <= 'z')
            {
                caps = !caps; // switch to caps
            }
        }
        // if capslock enable
        if (caps)
        {
            extension->colume = 1;
        }
        // data with e0
        if (extension->code_with_e0)
        {
            extension->colume = 2;
        }
        // get data form colume
        key = keyrow[extension->colume];

        // control key set
        switch (key)
        {
        case KEYBOARD_SHIFT_L:
            extension->l_shift = make;
            break;
        case KEYBOARD_SHIFT_R:
            extension->r_shift = make;
            break;
        case KEYBOARD_CTRL_L:
            extension->l_ctrl = make;
            break;
        case KEYBOARD_CTRL_R:
            extension->r_ctrl = make;
            break;
        case KEYBOARD_ALT_L:
            extension->l_alt = make;
            break;
        case KEYBOARD_ALT_R:
            extension->r_alt = make;
            break;
        case KEYBOARD_CAPS_LOCK:
            if (make) // set capslock status and set led when press CapsLock
            {
                extension->capslock = !extension->capslock;
                KeyboardSetLed(extension);
            }
            break;
        case KEYBOARD_NUM_LOCK:
            if (make)
            {
                extension->numlock = !extension->numlock;
                KeyboardSetLed(extension);
            }
            break;
        case KEYBOARD_SCROLL_LOCK:
            if (make)
            {
                extension->scrolllock = !extension->scrolllock;
                KeyboardSetLed(extension);
            }
            break;
        default:
            break;
        }

        // deal pad
        int pad = 0;
        if ((key >= KEYBOARD_PAD_MINUS) && (key <= KEYBOARD_PAD_9))
        {
            pad = 1;
            switch (key)
            {
            case KEYBOARD_PAD_SLASH:
                key = '/';
                break;
            case KEYBOARD_PAD_STAR:
                key = '*';
                break;
            case KEYBOARD_PAD_MINUS:
                key = '-';
                break;
            case KEYBOARD_PAD_PLUS:
                key = '+';
                break;
            case KEYBOARD_PAD_ENTER:
                key = KEYBOARD_ENTER;
                break;
            default:
                // number area
                if (extension->numlock && (key >= KEYBOARD_PAD_0) && (key <= KEYBOARD_PAD_9))
                {
                    key = (key - KEYBOARD_PAD_0) + '0';
                }
                else
                {
                    if (extension->numlock && (key == KEYBOARD_PAD_DOT))
                    {
                        key = '.';
                    }
                    else
                    {
                        switch (key)
                        {
                        case KEYBOARD_PAD_HOME:
                            key = KEYBOARD_HOME;
                            break;
                        case KEYBOARD_PAD_END:
                            key = KEYBOARD_END;
                            break;
                        case KEYBOARD_PAD_PAGEDOWN:
                            key = KEYBOARD_PAGEDOWN;
                            break;
                        case KEYBOARD_PAD_PAGEUP:
                            key = KEYBOARD_PAGEUP;
                            break;
                        case KEYBOARD_PAD_INS:
                            key = KEYBOARD_INSERT;
                            break;
                        case KEYBOARD_PAD_UP:
                            key = KEYBOARD_UP;
                            break;
                        case KEYBOARD_PAD_DOWN:
                            key = KEYBOARD_DOWN;
                            break;
                        case KEYBOARD_PAD_DELETE:
                            key = KEYBOARD_DELETE;
                            break;
                        default:
                            break;
                        }
                    }
                }
                break;
            }
        }
        // ctrl+alt+shift+key
        key |= extension->l_shift ? KEYBOARD_FLAG_SHIFT_L : 0;
        key |= extension->r_shift ? KEYBOARD_FLAG_SHIFT_R : 0;
        key |= extension->l_ctrl ? KEYBOARD_FLAG_CTRL_L : 0;
        key |= extension->r_ctrl ? KEYBOARD_FLAG_CTRL_R : 0;
        key |= extension->l_alt ? KEYBOARD_FLAG_ALT_L : 0;
        key |= extension->r_alt ? KEYBOARD_FLAG_ALT_R : 0;

        // set break flag
        key |= (make ? 0 : KEYBOARD_FLAG_BREAK);

        // set lock flag
        key |= extension->numlock ? KEYBOARD_FLAG_NUM : 0;
        key |= extension->capslock ? KEYBOARD_FLAG_CAPS : 0;
    }
    return key;
}

static int KeyboardProbe(device_extension_t *extension, uint8_t port)
{
    uint8_t type = 0xff;

    if (!extension)
        return -1;

    type = CheckDeviceType(port);
    if (type != PS2_TYPE_KEYBOARD)
        return -1;
    extension->type = type;
    extension->id = port;
    KPrint("[keyboard] found keyboard on port %d\n", port);
    return 0;
}

static void KeyboardSetLed(device_extension_t *extension)
{
    /*uint8_t leds = (extension->capslock << 2) | (extension->numlock << 1) | extension->scrolllock;
    // send set led cmd
    WriteDataTo(extension->id, KEYBOARD_CMD_SETLED);
    // write led data
    WriteDataTo(extension->id, leds);*/
}

// check output buffer
static int CheckOutBuffer()
{
    uint8_t flag, status;
    // read controller status
    status = In8(PS2CTRL_IO_STATUS_PORT);
    // test status flags
    if (status & PS2CTRL_STATUS_OUTFULL) // no is empty
        return 1;

    return 0;
}

static inline void CheckAck()
{
    uint8_t data;

    while ((data = ReadData()) != KEYBOARD_RET_ACK)
        ;
}

static int KeyboardHandler(irqno_t irq, void *data)
{
    device_extension_t *extension = (device_extension_t *)data;
    uint8_t scancode = In8(PS2CTRL_IO_DATA_PORT); // get key code

    if (extension->open) // put a keycode to fifo buffer
    {
        FifoIoPut(&extension->fifoio, scancode);
    }

    return 0;
}

static void KeyboardThread(void *arg)
{
    device_extension_t *extension = arg;
    uint32_t key;
    input_event_t event;

    while (1)
    {
        key = 0; // reset key
        key = KeyboardDoRead(extension);

        if (key > 0 && (key & KEYBOARD_KEY_MASK))
        {
            if (key & KEYBOARD_FLAG_BREAK)
            {
                event.value = 0;
            }
            else
            {
                event.value = 1;
            }
            event.type = EVENT_KEY; // key event
            event.code = ScanCodeToEventCode(key & KEYBOARD_KEY_MASK);
            InputEventPut(&extension->eventbuff, &event);
#ifdef DEBUG_KEYBOARD
            KPrint(PRINT_DEBUG "key event: type=%d code=%d value=%d\n", event.type, event.code, event.value);
            KPrint(PRINT_DEBUG "print key:%c-%x\n", key & KEYBOARD_KEY_MASK, key & KEYBOARD_KEY_MASK);
#endif
        }
    }
}

static iostatus_t KeyboardClose(device_object_t *device, io_request_t *ioreq)
{
    device_extension_t *extension = device->device_extension;
    iostatus_t status = IO_SUCCESS;

    if (!extension->exist)
        status = IO_FAILED;
    extension->open = 0;
    ioreq->io_status.status = IO_SUCCESS;
    ioreq->io_status.info = 0;
    IoCompleteRequest(ioreq);
    return status;
}

static iostatus_t KeyboardDevCtl(device_object_t *device, io_request_t *ioreq)
{
    device_extension_t *extension = device->device_extension;
    uint32_t cmd = ioreq->parame.devctl.code;
    uint32_t arg = ioreq->parame.devctl.arg;
    uint32_t led;

    switch (cmd)
    {
    case KEYBOARDIO_GETLED:
        led = extension->numlock | (extension->capslock << 1) |
              (extension->scrolllock << 2);
        *(uint32_t *)arg = led;
        break;
    case KEYBOARDIO_SETFLAGS:
        extension->flags = *(uint32_t *)arg;
        break;
    case KEYBOARDIO_GETFLAGS:
        *(uint32_t *)arg = extension->flags;
        break;
    default:
        break;
    }
}

static iostatus_t KeyboardRead(device_object_t *device, io_request_t *ioreq)
{
    device_extension_t *extension = device->device_extension;
    iostatus_t status = IO_SUCCESS;
    input_event_t *event = ioreq->user_buff;

    if (event && ioreq->parame.read.len == sizeof(input_event_t))
    {
        if (extension->open)
        {
            if (extension->flags & DEVICE_NOWAIT) // assign NOWAIT,only get once
            {
                if (InputEventGet(&extension->eventbuff, event))
                {
                    status = IO_FAILED;
                }
                else
                {
#ifdef DEBUG_DRIVER
                    //   KPrint("InputEventGet: type:%d code=%d value=%d\n", event->type, event->code, event->value);
#endif
                }
            }
            else
            {
                // read unit succeed
                while (1)
                {
                    if (!InputEventGet(&extension->eventbuff, event))
                        break;
                    TaskYield();
                }
#ifdef DEBUG_DRIVER
                // KPrint("InputEventGet: type:%d code=%d value=%d\n", event->type, event->code, event->value);
#endif
            }
        }
        else
        {
            status = IO_FAILED;
        }
    }
    else
    {
        status = IO_FAILED;
    }
    ioreq->io_status.status = status;
    ioreq->io_status.info = ioreq->parame.read.len; // return read length
    IoCompleteRequest(ioreq);
    return status;
}

iostatus_t KeyboardDriverFunc(driver_object_t *driver)
{
    iostatus_t status = IO_SUCCESS;

    // bind driver info
    driver->driver_enter = KeyboardEnter;
    driver->driver_exit = KeyboardExit;

    driver->dispatch_fun[IOREQ_OPEN] = KeyboardOpen;
    driver->dispatch_fun[IOREQ_CLOSE] = KeyboardClose;
    driver->dispatch_fun[IOREQ_READ] = KeyboardRead;
    driver->dispatch_fun[IOREQ_DEVCTL] = KeyboardDevCtl;
    // driver name
    string_new(&driver->name, DRIVER_NAME, DRIVER_NAME_LEN);

    return status;
}

static __init void Ps2KeyboardDriverEntry()
{
    KPrint("[driver] create keyboard driver\n");
    if (DriverObjectCreate(KeyboardDriverFunc) < 0)
    {
        KPrint(PRINT_ERR "[driver]ps2 keyboard driver create failed!\n");
    }
}

driver_initcall(Ps2KeyboardDriverEntry);